DLO-JZ Fully Sharded Data parallelism - Jour 4¶
Utilisation de la FSDP sur un modèle de langue Llama 3.2 3B.
Objet du notebook¶
Le but de ce notebook est d'optimiser un code d'apprentissage d'un modèle Llama 3.2 sur un dataset de roleplay Imagenet :
- Passage de DDP à FSDP
- FSDP optimisée grâce au wrap récursif
- Bonus : Application de la compilation PyTorch par dessus la FSDP
Les cellules dans ce notebook ne sont pas prévues pour être modifiées, sauf rares exceptions indiquées dans les commentaires. Les TP se feront en modifiant le code fsdp.py.
Les directives de modification seront marquées par l'étiquette TODO : dans le notebook suivant.
Les solutions sont présentes dans le répertoire solutions/.
Notebook rédigé par l'équipe assistance IA de l'IDRIS, octobre 2024
Environnement de calcul¶
Les fonctions python de gestion de queue Slurm dévelopées par l'IDRIS et les fonctions dédiées à la formation DLO-JZ sont à importer.
Le module d'environnement pour les jobs et la taille des images sont fixés pour ce notebook.
TODO : choisir un pseudonyme pour vous différencier dans la queue Slurm et dans les outils collaboratifs pendant la formation.
from idr_pytools import display_slurm_queue, gpu_jobs_submitter, search_log
from dlojz_tools import controle_technique, compare, GPU_underthehood, plot_accuracy, lrfind_plot, pipe_memory, turbo_profiler, comm_profiler
MODULE = 'pytorch-gpu/py3/2.4.0'
account = 'for@a100'
name = 'NATHOUNET'
Gestion de la queue Slurm¶
Cette partie permet d'afficher et de gérer la queue Slurm.
Pour afficher toute la queue utilisateur :
#display_slurm_queue(name)
Remarque: Cette fonction utilisée plusieurs fois dans ce notebook permet d'afficher la queue de manière dynamique, rafraichie toutes les 5 secondes. Cependant elle ne s'arrête que lorsque la queue est vide. Si vous désirez reprendre la main sur le notebook, il vous suffira d'arrêter manuellement la cellule avec le bouton stop. Cela a bien sûr aucun impact sur le scheduler Slurm. Les jobs ne seront pas arrêtés.
Si vous voulez arrêter des jobs dans la queue :
- Annuler tous vos jobs dans la queue (décommenter la ligne suivante)
- Annuler un job dans votre queue (décommenter la ligne suivante et ajouter le numéro du job à la fin de la ligne)
#!scancel -u $USER
Différence de scripts ¶
Pour le debug ou pour comparer son code avec les solutions mises à disposition, la fonction suivante permet d'afficher une page html contenant un différentiel de fichiers texte.
s1 = "./solutions/fsdp_2.py"
s2 = "./solutions/fsdp_3.py"
compare(s1, s2)
Voir le résultat du différentiel de fichiers sur la page suivante (attention au spoil !) :
Première exécution en DDP¶
!cp solutions/fsdp_0.py fsdp.py
Prenez connaissance du script fsdp.py. C'est un fine-tuning d'un modèle de langue Llama 3.2 à 3 milliards de paramètres sur un dataset de Roleplay récupéré sur HuggingFace.
La structure entre ce script et celui lié à la computer vision des autres jours est très similaire.
Note : on utilise les mêmes visualisations que dans les autres TPs donc les schémas peuvent parler d'images, mais dans notre contexte de NLP, la dimension du batch ne fait pas référence à des images mais plutôt aux nombres de séquences.
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 334759 jobid = ['334759']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
334759 gpu_p5 NATHOUNE cfor032 R 1:32 1 jean-zay-iam31
Done!
#jobid = ['334031']
controle_technique(jobid)
Train throughput: 28.75 images/second GPU throughput: 28.80 images/second epoch time: 200.37 seconds ----------- training step time average (fwd/bkwd on GPU): 0.555491 sec (38.9%/57.5%) +/- 0.053871 loading step time average (IO + CPU to GPU transfer): 0.001095 sec +/- 0.000453
turbo_profiler(jobid)
>>> Turbo Profiler >>> Training complete in 29.487146 s
comm_profiler(jobid, n_display=100)
Test d'occupation mémoire¶
Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.
Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.
Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.
n_gpu = 4
batch_sizes = [1, 2, 4, 8]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 335115 Submitted batch job 335120 Submitted batch job 335125 Submitted batch job 335132 jobids = ['335115', '335120', '335125', '335132']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
335115 gpu_p5 NATHOUNE cfor032 R 0:47 1 jean-zay-iam31
335120 gpu_p5 NATHOUNE cfor032 R 0:47 1 jean-zay-iam31
335125 gpu_p5 NATHOUNE cfor032 R 0:47 1 jean-zay-iam01
Done!
GPU_underthehood(jobids)
Batch size per GPU: 4 Max GPU Memory Allocated: 72.08 GB, Troughput: 10.616 images/second Batch size per GPU: 8 Max GPU Memory Allocated: 72.32 GB, Troughput: 18.526 images/second Batch size per GPU: 16 Max GPU Memory Allocated: 72.81 GB, Troughput: 28.875 images/second Batch size per GPU: 32 CUDA out of memory Memory occupancy by Model part : 71.835 +/- 0.003 GB
Passage au Fully Sharded Data Parallelism¶
TODO: Remplacez le Distributed Data Parallelism par le Fully Sharded Data Parallelism. Ce n'est qu'un simple wrapper et demande peu de modifications. Indice : ctrl-F de "#### Distribute the Model" pour repérer l'endroit où faire ça.
Important: Dans la documentation, vous pouvez voir un paramètre auto_wrap_policy, il est abordé dans la section suivante donc laissez le non précisé pour l'instant.
FSDP est un wrapper très haut niveau qui fait toutes les communications de manière cachée pour faciliter son utilisation. C'est son grand avantage par rapport à DeepSpeed.
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00')
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336009 jobid = ['336009']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336009 gpu_p5 NATHOUNE cfor032 R 0:44 1 jean-zay-iam31
Done!
#jobid = ['902284']
controle_technique(jobid)
Train throughput: 32.89 images/second GPU throughput: 32.91 images/second epoch time: 175.12 seconds ----------- training step time average (fwd/bkwd on GPU): 0.486194 sec (60.1%/33.0%) +/- 0.065956 loading step time average (IO + CPU to GPU transfer): 0.000251 sec +/- 0.000032
pipe_memory(jobid)
comm_profiler(jobid)
Test d'occupation mémoire¶
Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.
Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.
Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.
n_gpu = 4
batch_sizes = [2, 4, 8, 16]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336026 Submitted batch job 336027 Submitted batch job 336028 Submitted batch job 336029 jobids = ['336026', '336027', '336028', '336029']
#jobids = ['299137', '299138', '299139', '299140', '299141']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336028 gpu_p5 NATHOUNE cfor032 CG 1:10 1 jean-zay-iam01
Done!
GPU_underthehood(jobids)
Batch size per GPU: 8 Max GPU Memory Allocated: 45.39 GB, Troughput: 22.542 images/second Batch size per GPU: 16 Max GPU Memory Allocated: 45.88 GB, Troughput: 32.915 images/second Batch size per GPU: 32 Max GPU Memory Allocated: 53.17 GB, Troughput: 42.554 images/second Batch size per GPU: 64 CUDA out of memory Memory occupancy by Model part : 41.743 +/- 3.158 GB
Contrôle technique de la configuration optimale¶
controle_technique([jobids[-2]])
Train throughput: 42.51 images/second GPU throughput: 42.55 images/second epoch time: 135.50 seconds ----------- training step time average (fwd/bkwd on GPU): 0.751985 sec (54.4%/40.4%) +/- 0.078437 loading step time average (IO + CPU to GPU transfer): 0.000791 sec +/- 0.000241
Optimisation de la FSDP en l'appliquant récursivement¶
En utilisant la FSDP sans l'argument auto_wrap_policy, alors tous les paramètres du modèle sont transmis en une seule fois plutôt que quand ils sont nécessaires. On pourrait imaginer transmettre d'abord les paramètres des premières couches (dont on a besoin rapidement), puis ensuite les autres.
Le paramètre auto_wrap_policy permet de faire ça. Il est puissant car il permet de maximiser le recouvrement entre les calculs et les communications et minimise la mémoire nécessaire.
TODO: Ajoutez le paramètre auto_wrap_policy. Indice : Pour les transformers, il est de coutume de wrapper chaque block de transformer indépendamment (signature de la méthode transformer_auto_wrap_policy ici).
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336276 jobid = ['336276']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336276 gpu_p5 NATHOUNE cfor032 R 0:45 1 jean-zay-iam31
Done!
#jobid = ['902284']
controle_technique(jobid)
Train throughput: 34.64 images/second GPU throughput: 34.65 images/second epoch time: 166.30 seconds ----------- training step time average (fwd/bkwd on GPU): 0.461701 sec (43.9%/53.1%) +/- 0.028161 loading step time average (IO + CPU to GPU transfer): 0.000237 sec +/- 0.000031
comm_profiler(jobid, n_display=100)
Vous pouvez consulter les logs de l'entraînement. On affiche le modèle, et on peut par conséquent voir l'impact de l'argument auto_wrap_policy.
Test d'occupation mémoire¶
Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.
Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.
Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.
n_gpu = 4
batch_sizes = [4, 8, 12, 16, 32]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336313 Submitted batch job 336314 Submitted batch job 336315 Submitted batch job 336316 Submitted batch job 336318 jobids = ['336313', '336314', '336315', '336316', '336318']
#jobids = ['300073', '300074', '300075', '300076', '300077']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336316 gpu_p5 NATHOUNE cfor032 R 2:18 1 jean-zay-iam02
Done!
GPU_underthehood(jobids)
Batch size per GPU: 16 Max GPU Memory Allocated: 29.56 GB, Troughput: 34.349 images/second Batch size per GPU: 32 Max GPU Memory Allocated: 42.67 GB, Troughput: 43.954 images/second Batch size per GPU: 48 Max GPU Memory Allocated: 55.78 GB, Troughput: 48.743 images/second Batch size per GPU: 64 Max GPU Memory Allocated: 68.89 GB, Troughput: 29.467 images/second Batch size per GPU: 128 CUDA out of memory Memory occupancy by Model part : 26.281 +/- 9.830 GB
Contrôle technique de la configuration optimale¶
controle_technique(jobids[-3])
Train throughput: 48.73 images/second GPU throughput: 48.74 images/second epoch time: 118.20 seconds ----------- training step time average (fwd/bkwd on GPU): 0.984761 sec (41.3%/56.5%) +/- 0.043922 loading step time average (IO + CPU to GPU transfer): 0.000269 sec +/- 0.000037
Bonus : torch.compile par dessus FSDP¶
TODO: Appliquez la compilation par PyTorch de votre modèle.
Indice: ctrl-F de "#### JIT" pour trouver où faire ça.
n_gpu = 4
command = f'fsdp.py --batch-size 16 --num-workers 2 --seq-len 512 --test --compile --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336937 jobid = ['336937']
#jobid = ['902284']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336937 gpu_p5 NATHOUNE cfor032 R 2:37 1 jean-zay-iam02
Done!
controle_technique(jobid)
Train throughput: 64.93 images/second GPU throughput: 64.95 images/second epoch time: 88.71 seconds ----------- training step time average (fwd/bkwd on GPU): 0.985422 sec (40.5%/57.3%) +/- 0.042639 loading step time average (IO + CPU to GPU transfer): 0.000297 sec +/- 0.000049
comm_profiler(jobid, n_display=100)
Test d'occupation mémoire¶
Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.
Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.
Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.
n_gpu = 4
batch_sizes = [8, 12, 16, 20, 24]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test --compile'
for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task Submitted batch job 336987 Submitted batch job 336989 Submitted batch job 336991 Submitted batch job 336994 Submitted batch job 336995 jobids = ['336987', '336989', '336991', '336994', '336995']
#jobids = ['300403', '300404', '300406', '300407', '300408', '300409']
display_slurm_queue(name)
JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
336994 gpu_p5 NATHOUNE cfor032 R 2:48 1 jean-zay-iam31
Done!
GPU_underthehood(jobids)
Batch size per GPU: 32 Max GPU Memory Allocated: 35.62 GB, Troughput: 55.736 images/second Batch size per GPU: 48 Max GPU Memory Allocated: 45.21 GB, Troughput: 61.295 images/second Batch size per GPU: 64 Max GPU Memory Allocated: 54.80 GB, Troughput: 64.816 images/second Batch size per GPU: 80 Max GPU Memory Allocated: 64.38 GB, Troughput: 76.294 images/second Batch size per GPU: 96 CUDA out of memory Memory occupancy by Model part : 36.420 +/- 5.592 GB
Contrôle technique de la configuration optimale¶
controle_technique(jobids[-2])
Train throughput: 69.96 images/second GPU throughput: 76.29 images/second epoch time: 82.33 seconds ----------- training step time average (fwd/bkwd on GPU): 1.048578 sec (36.0%/64.8%) +/- 0.065510 loading step time average (IO + CPU to GPU transfer): 0.094912 sec +/- 0.047183
torch.compile est encore très nouveau et il peut arriver qu'un modèle ne puisse pas être converti. Plusieurs backends sont disponibles (voir documentation officielle). Dans les cas des modèles les plus exotiques, la compilation peut tout simplement échoué. C'est tellement bas niveau qu'il est bien possible qu'on ne puisse rien y faire, c'est juste lié au fait que torch.compile est relativement nouveau. À garder à l'esprit cependant, car cela peut augmenter de 50%, voire parfois 100% le throughput de votre modèle.
